#!/usr/bin/python
# -*- coding: UTF-8 -*-
import base64
import io

import qianfan
import os

from PIL import Image


class AIModel(object):
    __slots__ = ('model', '__chat_comp')

    def __init__(self, model):
        self.model = model
        os.environ["QIANFAN_ACCESS_KEY"] = "ALTAKVMW05mav4NilQNwpHDlSl"
        os.environ["QIANFAN_SECRET_KEY"] = "0030dc450f60423caeef2b1d390a9572"
        self.__chat_comp = qianfan.ChatCompletion()

    def replay(self, user_msg: str) -> str:
        resp = self.__chat_comp.do(model=self.model, stop=['。'], messages=[{
            'role': 'user',
            'content': user_msg
        }])

        return resp['body']['result']


class DrawModel(object):
    __slots__ = 't2i'

    def __init__(self):
        os.environ["QIANFAN_ACCESS_KEY"] = "ALTAKVMW05mav4NilQNwpHDlSl"
        os.environ["QIANFAN_SECRET_KEY"] = "0030dc450f60423caeef2b1d390a9572"
        self.t2i = qianfan.Text2Image()

    def reply(self, keywords):
        resp = self.t2i.do(prompt=keywords, with_decode="base64")["body"]["data"][0]["b64_image"]
        return base64.b64decode(resp)


if __name__ == '__main__':
    dr = DrawModel()
    print(dr.reply('2'))